import math
import sys
sys.path.append('../ste/')
from core import *
import matplotlib.pyplot as plt
import numpy as np
from visualizer.basic import *
from sklearn import mixture
from sklearn.cluster import KMeans
import copy

def Ts(ts2, ts1, l, m, step, func):
	return TE(ts2, ts1, l, m, step, func) - TE(ts1, ts2, l, m, step, func)


def gTE(ts1, ts2, l, m, step=1):
	N = len(ts1)
	x = []
	y = []
	for i in range(N - l * (m - 1)):
		subsequence1 = copy.deepcopy(ts1[i:i + l * (m - 1) + 1:l])
		subsequence2 = copy.deepcopy(ts2[i:i + l * (m - 1) + 1:l])
		subsequence1 -= np.min(subsequence1)
		subsequence2 -= np.min(subsequence2)
		if np.max(subsequence1) != 0:
			subsequence1 /= np.max(subsequence1)
		else:
			subsequence1 = np.zeros(len(subsequence1))
		if np.max(subsequence2) != 0:
			subsequence2 /= np.max(subsequence2)
		else:
			subsequence2 = np.zeros(len(subsequence2))
		x.append(subsequence1)
		y.append(subsequence2)
	
	kmeans = KMeans(n_clusters=math.factorial(m) if math.factorial(m)<100 else 100).fit(y)
	kmeans.fit(x)
	x = kmeans.predict(x)
	kmeans.fit(y)
	y = kmeans.predict(y)

	sequence_x1 = x[step:]
	sequence_x0 = x[:-step]
	sequence_y0 = y[:-step]
	Nx1x0y0 = defaultdict(int)
	Nx1x0 = defaultdict(int)
	Nx0y0 = defaultdict(int)
	Nx0 = defaultdict(int)
	N = len(sequence_x1)

	for i in range(N):
		Nx1x0y0[(sequence_x1[i], sequence_x0[i], sequence_y0[i])] += 1
		Nx1x0[(sequence_x1[i], sequence_x0[i])] += 1
		Nx0y0[(sequence_x0[i], sequence_y0[i])] += 1
		Nx0[sequence_x0[i]] += 1

	# print(Nx1x0y0, Nx1x0, Nx0y0, Nx0)

	px1dx0y0 = defaultdict(int)
	px1dx0 = defaultdict(int)
	px1x0y0 = defaultdict(int)

	for (x1, x0, y0) in Nx1x0y0.keys():
		if Nx1x0[(x1, x0)] != 0:
			px1dx0y0[(x1, x0, y0)] = Nx1x0y0[(x1, x0, y0)] / Nx0y0[(x0, y0)]
		if Nx0[x0] != 0:
			px1dx0[(x1, x0)] = Nx1x0[(x1, x0)] / Nx0[x0]
		if Nx0y0[(x0, y0)] != 0:
			px1x0y0[(x1, x0, y0)] = Nx1x0y0[(x1, x0, y0)] / N

	# print('N calculation done')
	Tyx = 0
	# print(px1dx0y0, px1x0y0, px1dx0)
	for (x1, x0, y0) in Nx1x0y0.keys():
		if px1dx0[(x1, x0)] != 0 and px1x0y0[(x1, x0, y0)] != 0:
			Tyx += px1x0y0[(x1, x0, y0)] * np.log2(px1dx0y0[(x1, x0, y0)] / px1dx0[(x1, x0)])
	
	return Tyx

if __name__ == '__main__':
	print('ing...')
	ts1 = np.loadtxt('../data/pr_NY.csv')  # 降水量 # Y
	ts2 = np.loadtxt('../data/tas_NY.csv')  # 温度 # X
	N = 21
	l = 1
	m = 3
	step = 1
	# symbol
	symbol_0 = []
	x = list(range(3, N))
	for m in range(3, N):
		T = Ts(ts2, ts1, l, m, step, symbol)
		symbol_0.append(T)
	
	T_series = []
	for m in range(3, N):
		T = gTE(ts2, ts1, l, m, step) - gTE(ts1, ts2, l, m, step)
		T_series.append(T)
	plt.plot(x, symbol_0, label='STE')
	plt.plot(x, T_series, label='Kmeans STE')
	plt.xlabel('$m$')
	plt.ylabel('$T(m)$')
	plt.xticks(np.arange(3, N, 3), np.arange(3, N, 3))
	plt.legend()
	plt.tight_layout()
	plt.savefig(PATH + 'clusterTE2.pdf', format='pdf')
	plt.show()
	print('done')